# encoding: utf-8
# author: 侯佳涛
# file_name: data_loader
# create_time: 2021/2/26, 下午2:53
import codecs
import os

import data_utils


def load_sentence(data_path):
    """
    加载数据
    :param data_path: 数据路径
    :return: [[[s1_w1, s1_t1], [s1_w2, s1_t2], ...], [[s2_w1, s2_t1], [s2_w2, s2_t2], ...], ...]
    """
    sentences = []
    sentence = []

    with codecs.open(data_path, 'r', encoding="UTF-8") as f:
        data = f.readlines()

    for line in data:
        line = line.rstrip()
        if not line:  # 当前句子结束
            sentences.append(sentence)
            sentence = []
            continue
        else:
            if line[0] == ' ':  # 说明是'  tag'
                sentence.append(['$', line.strip()])
            else:
                sentence.append(line.split())

    if sentence:
        sentences.append(sentence)

    return sentences


def check_tags(sentences):
    """
    检测编码
    :param sentences:
    :return:
    """
    tags = [[w_t[-1] for w_t in sentence] for sentence in sentences]

    data_utils.check_bio(tags)

    for i, tag_list in enumerate(tags):
        for j, tag in enumerate(tag_list):
            sentences[i][j][-1] = tag

    return sentences


def tag_mapping(sentences):
    """
    构建标签映射
    :param sentences:
    :return:
    """
    tags = [[w_t[-1] for w_t in sentence] for sentence in sentences]

    return data_utils.create_mapping(tags)


def prepare_dataset(sentences, mode, params, tag2id):
    """
    将数据处理成标准格式:
        - string
        - ids
        - mask
        - segment_ids
        - label_ids
    :param sentences:
    :param mode: train, dev, test
    :param params:
    :param tag2id
    :return:
    """
    assert mode.lower in ['train', 'dev', 'test']
    return_data = []

    for sentence in sentences:
        sample = data_utils.convert_sentence_to_sample(
            sentence,
            tag2id,
            params
        )

        return_data.append(sample)

    return return_data


if __name__ == '__main__':
    sentences = load_sentence("./data/example.test")
    sentences = check_tags(sentences)
    # word2id, id2word = char_mapping(sentences)
    tag2id, id2tag = tag_mapping(sentences)
    print(tag2id, id2tag)
